import iotdb_api
import random
import numpy as np
def create_database(database_name):
    """
    在 IotDB 中创建数据库
    :param database_name: 数据库名
    :return: 操作状态
    """

    create_database = "create database root." + database_name
    iotdb_api.non_query(create_database)


def create_timeseries(database_name, timeseries_name):
    """
    在当前数据库下建立时间序列
    :param database_name: 数据库名
    :param timeseries_name: 时间序列名
    :return: 操作状态
    """

    create_timeseries = "create timeseries root." + database_name + "." + timeseries_name + " TEXT encoding=PLAIN"
    iotdb_api.non_query(create_timeseries)


def creeate():
    """
    创建数据库以及时序序列
    :return:
    """

    # 创建数据库
    # root.all_msg，记录流量包内的所有信息
    create_database("all_msg")
    # root.company，记录公司
    create_database("company")
    # 创建数据库 root.province
    create_database("province")

    # 创建 root.all_msg 数据库下的时序序列
    # root.all_msg.iden_id 记录随流标识 ID
    create_timeseries("all_msg", "iden_id")
    # root.all_msg.iden_sum 记录随流标识 ID 摘要
    create_timeseries("all_msg", "iden_sum")
    # root.all_msg.ser_id 记录业务数据 ID
    create_timeseries("all_msg", "ser_id")
    # root.all_msg.data_size 记录已发送数据量
    create_timeseries("all_msg", "data_size")
    # root.all_msg.src_id 记录源企业 ID
    create_timeseries("all_msg", "src_id")
    # root.all_msg.dst_id 记录目标企业 ID
    create_timeseries("all_msg", "dst_id")
    # root.all_msg.level 记录敏感等级
    create_timeseries("all_msg", "level")
    # root.all_msg.model 记录数据出境模式
    create_timeseries("all_msg", "model")
    # root.all_msg.pro_type 记录协议类型
    create_timeseries("all_msg", "pro_type")
    # root.all_msg.src_ip 记录源企业 IP
    create_timeseries("all_msg", "src_ip")
    # root.all_msg.dst_ip 记录目的企业 IP
    create_timeseries("all_msg", "dst_ip")
    create_timeseries("all_msg", "nation")
    create_timeseries("all_msg", "route")

    # 创建 root.company 数据库下的时序序列
    # root.company.src_id 记录公司编号
    create_timeseries("company", "src_id")
    # root.company.size
    create_timeseries("company", "data_size")
    # root.company.model
    create_timeseries("company", "level")

    # 创建 root.province 数据库下的时间序列
    # root.province.id 记录省份信息
    create_timeseries("province", "prov_id")
    # root.province.size 记录流经该省份的数据大小
    create_timeseries("province", "data_size")
    # root.province.dst_ip 记录该省份出口流量
    create_timeseries("province", "dst_ip")

provinces = [  
    "北京市", "天津市", "上海市", "重庆市",  
    "河北省", "山西省", "辽宁省", "吉林省", "黑龙江省",  
    "江苏省", "浙江省", "安徽省", "福建省", "江西省",  
    "山东省", "河南省", "湖北省", "湖南省", "广东省",  
    "海南省", "四川省", "贵州省", "云南省", "陕西省",  
    "甘肃省", "青海省", "内蒙古自治区",  
    "广西壮族自治区", "西藏自治区", "宁夏回族自治区",  
    "新疆维吾尔自治区",  
]  

cities_in = [
    "天津市", "重庆市",  
    "哈尔滨市", "长春市", "沈阳市", "石家庄市",  
    "呼和浩特市", "太原市", "济南市", "郑州市",  
    "西安市", "兰州市", "银川市", "西宁市",  
    "乌鲁木齐市", "南京市", "杭州市", "合肥市",  
    "福州市", "南昌市", "长沙市", "武汉市",  
    "南宁市", "海口市",  
    "成都市", "贵阳市", "昆明市", "拉萨市",  
]
cities_out = [
    "北京市","上海市","广州市"
]
cities = [
    "天津市", "重庆市",  
    "哈尔滨市", "长春市", "沈阳市", "石家庄市",  
    "呼和浩特市", "太原市", "济南市", "郑州市",  
    "西安市", "兰州市", "银川市", "西宁市",  
    "乌鲁木齐市", "南京市", "杭州市", "合肥市",  
    "福州市", "南昌市", "长沙市", "武汉市",  
    "南宁市", "海口市",  
    "成都市", "贵阳市", "昆明市", "拉萨市",  
    "北京市", "上海市","广州市"
]
provincial_capitals = {
    "北京市":"北京市", 
    "天津市":"天津市", 
    "上海市":"上海市", 
    "重庆市":"重庆市",  
    "河北省":"石家庄市",
    "山西省":"太原市", 
    "辽宁省":"沈阳市", 
    "吉林省":"长春市", 
    "黑龙江省":"哈尔滨市",  
    "江苏省":"南京市", 
    "浙江省":"杭州市", 
    "安徽省":"合肥市", 
    "福建省":"福州市", 
    "江西省":"南昌市",  
    "山东省":"济南市", 
    "河南省":"郑州市", 
    "湖北省":"武汉市", 
    "湖南省":"长沙市", 
    "广东省":"广州市",  
    "海南省":"海南市", 
    "四川省":"成都市", 
    "贵州省":"贵阳市", 
    "云南省":"昆明市", 
    "陕西省":"西安市",  
    "甘肃省":"兰州市", 
    "青海省":"西宁市", 
    "内蒙古自治区":"呼和浩特市",  
    "广西壮族自治区":"南宁市", 
    "西藏自治区":"拉萨市", 
    "宁夏回族自治区":"银川市",  
    "新疆维吾尔自治区":"乌鲁木齐市"  
}

nations= [
    "英国","法国","德国","意大利","西班牙","葡萄牙","美国","加拿大","朝鲜","韩国","日本","蒙古","俄罗斯","巴西","越南","老挝","柬埔寨","泰国","澳大利亚","新西兰","巴西"
]

id = 0
def example():
    """
    案例
    :return:
    """
    global id
    id += 1
    iden_id = str(id)
    iden_sum = "54544135"
    ser_id = str(random.randint(1,100))
    data_size = str(random.randint(1,200))+" MB"
    src_id = str(random.randint(1,200))
    dst_id = str(random.randint(1,200))
    level = str(random.randint(0,3))
    model = str(random.randint(1,15))
    pro_type = "ipv4"
    src_ip = "127.0.0.1"
    dst_ip = "127.0.0.5"
    nation = nations[random.randint(0,len(nations)-1)]
    prov_id = provinces[random.randint(0,len(provinces)-1)]
    subsidiary =random.randint(1,5)
    visited = np.zeros(len(cities))
    route_list =[]
    route_list.append(provincial_capitals[prov_id])
    in_num = 0
    while route_list[len(route_list)-1] not in cities_out :
        city = random.choice(cities)
        if city in cities_out:
            route_list.append(city)
            break
        elif visited[cities.index(city)] == 1 or in_num >= 5:
            route_list.append(random.choice(cities_out))
            break
        else:
            route_list.append(city)
            visited[cities.index(city)] 
            in_num += 1
    route_list.append(nation)
    route = '.'.join(route_list)
    iotdb_api.insert_all(iden_id, iden_sum, ser_id, data_size, src_id, dst_id, level, model, pro_type, src_ip, dst_ip, nation, route)
    iotdb_api.insert_company(src_id, data_size, level, model,subsidiary)
    iotdb_api.insert_province(prov_id, data_size, dst_ip)

def example2():
    """
    案例
    :return:
    """
    global id
    id += 1
    company_number = [3, 18, 70, 112, 136, 181]
    iden_id = str(id)
    iden_sum = "54544135"
    ser_id = str(random.randint(1,100))
    data_size = str(random.randint(1,200))+" MB"
    src_id = str(random.choice(company_number))
    dst_id = str(random.randint(1,200))
    level = str(random.randint(0,3))
    model = str(random.randint(1,15))
    pro_type = "ipv4"
    src_ip = "127.0.0.1"
    dst_ip = "127.0.0.5"
    nation = nations[random.randint(0,len(nations)-1)]
    prov_id = provinces[random.randint(0,len(provinces)-1)]
    subsidiary =random.randint(1,5)
    visited = np.zeros(len(cities))
    route_list =[]
    route_list.append(provincial_capitals[prov_id])
    in_num = 0
    while route_list[len(route_list)-1] not in cities_out :
        city = random.choice(cities)
        if city in cities_out:
            route_list.append(city)
            break
        elif visited[cities.index(city)] == 1 or in_num >= 5:
            route_list.append(random.choice(cities_out))
            break
        else:
            route_list.append(city)
            visited[cities.index(city)] 
            in_num += 1
    route_list.append(nation)
    route = '.'.join(route_list)
    iotdb_api.insert_all(iden_id, iden_sum, ser_id, data_size, src_id, dst_id, level, model, pro_type, src_ip, dst_ip, nation, route)
    iotdb_api.insert_company(src_id, data_size, level, model,subsidiary)
    iotdb_api.insert_province(prov_id, data_size, dst_ip)

if __name__ == '__main__':
    creeate()
    for i in range(200):
        example()
    for i in range(800):
        example2() #产生几个比例高的公司